package balancer

import (
	"sync"
	"sync/atomic"

	"go-gateworker/internal/config"
	"go-gateworker/internal/health"
	"go-gateworker/internal/models"
)

// Strategy 负载均衡策略
type Strategy int

const (
	StrategyRoundRobin Strategy = iota
	StrategyLeastConn
	StrategyRandom
)

// LoadBalancer 负载均衡器
type LoadBalancer struct {
	servers       []*models.Server
	strategy      Strategy
	threshold     int64
	roundRobin    uint64
	healthChecker *health.Checker
	mu            sync.RWMutex
}

// Config 负载均衡器配置
type Config struct {
	Strategy    Strategy
	Threshold   int64
	HealthCheck config.HealthCheckConfig
}

// New 创建新的负载均衡器
func New(servers []*models.Server, cfg Config) *LoadBalancer {
	lb := &LoadBalancer{
		servers:   servers,
		strategy:  cfg.Strategy,
		threshold: cfg.Threshold,
	}

	// 创建并启动健康检查器
	lb.healthChecker = health.NewChecker(servers, cfg.HealthCheck)
	go lb.healthChecker.Start()

	// 监听服务恢复通知
	go lb.listenForRecovery()

	return lb
}

// listenForRecovery 监听服务恢复通知
func (lb *LoadBalancer) listenForRecovery() {
	recoveryChannel := lb.healthChecker.GetRecoveryChannel()
	for server := range recoveryChannel {
		lb.mu.Lock()
		// 服务恢复后，更新服务器状态
		for _, s := range lb.servers {
			if s.URL.String() == server.URL.String() {
				s.MarkSuccess()
				break
			}
		}
		lb.mu.Unlock()
	}
}

// Stop 停止负载均衡器
func (lb *LoadBalancer) Stop() {
	if lb.healthChecker != nil {
		lb.healthChecker.Stop()
	}
}

// ForceHealthCheck 强制执行一次健康检查
func (lb *LoadBalancer) ForceHealthCheck() {
	if lb.healthChecker != nil {
		lb.healthChecker.ForceCheck()
	}
}

// GetServer 获取可用服务器
func (lb *LoadBalancer) GetServer() *models.Server {
	switch lb.strategy {
	case StrategyLeastConn:
		return lb.getLeastConnServer()
	case StrategyRoundRobin:
		return lb.getRoundRobinServer()
	default:
		return lb.getLeastConnServer()
	}
}

// getLeastConnServer 获取最小连接数的服务器
func (lb *LoadBalancer) getLeastConnServer() *models.Server {
	var bestServer *models.Server
	var minLoad int64 = -1

	// 优先选择主服务器（如果健康且未过载）
	if len(lb.servers) > 0 {
		mainServer := lb.servers[0]
		if mainServer.IsHealthy() {
			load := mainServer.GetLoad()
			if load < lb.threshold {
				return mainServer
			}
		}
	}

	// 在所有服务器中找到最佳可用服务器
	for _, server := range lb.servers {
		// 跳过不可用的服务器
		if !server.IsAvailable() {
			continue
		}

		load := server.GetLoad()
		if bestServer == nil || load < minLoad {
			minLoad = load
			bestServer = server
		}
	}

	return bestServer
}

// getRoundRobinServer 轮询方式获取服务器
func (lb *LoadBalancer) getRoundRobinServer() *models.Server {
	count := uint64(len(lb.servers))
	if count == 0 {
		return nil
	}

	// 尝试找到一个可用的服务器
	for i := uint64(0); i < count; i++ {
		next := atomic.AddUint64(&lb.roundRobin, 1) % count
		server := lb.servers[next]

		if server.IsAvailable() {
			return server
		}
	}
	return nil
}

// UpdateServers 更新服务器列表
func (lb *LoadBalancer) UpdateServers(servers []*models.Server) {
	lb.mu.Lock()
	defer lb.mu.Unlock()
	lb.servers = servers

	// 更新健康检查器的服务器列表
	if lb.healthChecker != nil {
		lb.healthChecker.UpdateServers(servers)
	}
}

// GetServerCount 获取服务器总数
func (lb *LoadBalancer) GetServerCount() int {
	return len(lb.servers)
}

// GetHealthyServerCount 获取健康的服务器数量
func (lb *LoadBalancer) GetHealthyServerCount() int {
	count := 0
	for _, server := range lb.servers {
		if server.IsHealthy() {
			count++
		}
	}
	return count
}

// GetAvailableServerCount 获取可用服务器数量（包括健康和过载但仍在运行的服务器）
func (lb *LoadBalancer) GetAvailableServerCount() int {
	count := 0
	for _, server := range lb.servers {
		if server.IsAvailable() {
			count++
		}
	}
	return count
}
